package com.master.chat.gpt.service.impl;

import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.elasticsearch.core.SearchResponse;
import co.elastic.clients.elasticsearch.core.search.Hit;
import co.elastic.clients.json.jackson.JacksonJsonpMapper;
import co.elastic.clients.transport.ElasticsearchTransport;
import co.elastic.clients.transport.rest_client.RestClientTransport;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.master.chat.api.base.config.LocalCache;
import com.master.chat.api.base.enums.ChatContentEnum;
import com.master.chat.api.base.enums.ChatModelEnum;
import com.master.chat.api.base.enums.ChatRoleEnum;
import com.master.chat.api.moonshot.MoonshotClient;
import com.master.chat.api.openai.OpenAiStreamClient;
import com.master.chat.api.openai.entity.chat.OpenAiMessage;
import com.master.chat.api.qianwen.QianWenClient;
import com.master.chat.api.wenxin.WenXinClient;
import com.master.chat.api.xfyun.SparkClient;
import com.master.chat.api.zhipu.ZhiPuClient;
import com.master.chat.common.enums.StatusEnum;
import com.master.chat.common.exception.BusinessException;
import com.master.chat.common.exception.ErrorException;
import com.master.chat.common.validator.ValidatorUtil;
import com.master.chat.framework.security.UserDetail;
import com.master.chat.gpt.enums.ChatStatusEnum;
import com.master.chat.gpt.pojo.command.ChatMessageCommand;
import com.master.chat.gpt.pojo.dto.ChatMessageDTO;
import com.master.chat.gpt.pojo.entity.User;
import com.master.chat.gpt.pojo.vo.ModelVO;
import com.master.chat.gpt.service.IChatMessageService;
import com.master.chat.gpt.service.IGptService;
import com.master.chat.gpt.service.IUserService;
import com.master.chat.gpt.service.SseService;
import com.master.chat.utils.SqlProcessor;
import com.zhipu.oapi.Constants;
import com.zhipu.oapi.service.v4.model.*;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.apache.http.HttpHost;
import org.apache.http.message.BasicHeader;
import org.elasticsearch.client.RestClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;

@Slf4j
@Service
public class AssistantChatServiceImpl implements SseService {


    private static final String[] drawingWords = {"画画", "作画", "画图", "绘画", "描绘"};
    private static final String[] drawingInstructions = {"请画", "画一", "画个",};

    private static final String ES_HOST = "http://localhost:9200";
    private static OpenAiStreamClient openAiStreamClient;
    private static WenXinClient wenXinClient;
    private static ZhiPuClient zhiPuClient;
    private static QianWenClient qianWenClient;
    private static SparkClient sparkClient;
    private static MoonshotClient moonshotClient;
    private final IGptService gptService;

    private final IUserService userService;

    @Autowired
    public AssistantChatServiceImpl(IGptService gptService, OpenAiStreamClient openAiStreamClient, WenXinClient wenXinClient,
                                    ZhiPuClient zhiPuClient, QianWenClient qianWenClient, SparkClient sparkClient, MoonshotClient moonshotClient,
                                    IUserService userService,
                                    SqlProcessor processor
    ) {
        this.gptService = gptService;
        this.userService = userService;
        AssistantChatServiceImpl.openAiStreamClient = openAiStreamClient;
        AssistantChatServiceImpl.wenXinClient = wenXinClient;
        AssistantChatServiceImpl.zhiPuClient = zhiPuClient;
        AssistantChatServiceImpl.qianWenClient = qianWenClient;
        AssistantChatServiceImpl.sparkClient = sparkClient;
        AssistantChatServiceImpl.moonshotClient = moonshotClient;
        AssistantChatServiceImpl.processor = processor;

    }

    private static SqlProcessor processor;

    @Override
    public SseEmitter createSse(String uid) {
        //默认30秒超时,设置为0L则永不超时
        SseEmitter sseEmitter = new SseEmitter(0L);
        //完成后回调
        sseEmitter.onCompletion(() -> {
            log.info("[{}]结束连接", uid);
            LocalCache.CACHE.remove(uid);
        });
        //超时回调
        sseEmitter.onTimeout(() -> {
            log.info("[{}]连接超时", uid);
        });
        //异常回调
        sseEmitter.onError(
                throwable -> {
                    try {
                        log.info("[{}]连接异常,{}", uid, throwable.toString());
                        sseEmitter.send(SseEmitter.event()
                                .id(uid)
                                .name("发生异常！")
                                .data(OpenAiMessage.builder().content("发生异常请重试！").build())
                                .reconnectTime(3000));
                        LocalCache.CACHE.put(uid, sseEmitter);
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                }
        );
        try {
            sseEmitter.send(SseEmitter.event().reconnectTime(5000));
        } catch (IOException e) {
            e.printStackTrace();
        }
        LocalCache.CACHE.put(uid, sseEmitter);
        log.info("[{}]创建sse连接成功！", uid);
        return sseEmitter;
    }

    @Override
    public void closeSse(String uid) {
        SseEmitter sse = (SseEmitter) LocalCache.CACHE.get(uid);
        if (sse != null) {
            sse.complete();
            //移除
            LocalCache.CACHE.remove(uid);
        }
    }


    @Override
    public void sseChat(UserDetail user, String conversationId, HttpServletResponse response) {
        ChatMessageDTO chatMessage = gptService.getMessageByConverstationId(conversationId);
        ModelVO model = gptService.getModel(chatMessage.getModel()).getData();
        if (ValidatorUtil.isNull(model)) {
            throw new BusinessException("模型已经不存在啦，请切换模型进行回复～");
        }
        if (StatusEnum.DISABLED.getValue().equals(model.getStatus())) {
            throw new BusinessException("该模型已被禁用，请切换模型进行回复～");
        }
        String prompt = chatMessage.getContent();
        String version = model.getVersion();
        ChatModelEnum modelEnum = ChatModelEnum.getEnum(chatMessage.getModel());

        String sysPrompt = model.getSysprompt();
        String ansSysPrompt = model.getAnssysprompt();


        SseEmitter sseEmitter = createSse(user.getUid());
        if (sseEmitter == null) {
            log.info("聊天消息推送失败uid:[{}],没有创建连接，请重试。", user.getUid());
            throw new BusinessException("聊天消息推送失败uid:[{}],没有创建连接，请重试。~");
        }
        List<ChatMessageDTO> chatMessages = gptService.listMessageByConverstationId(user.getId(), conversationId);
        if (ValidatorUtil.isNullIncludeArray(chatMessages)) {
            throw new BusinessException("消息发送失败");
        }
        // ChatGPT、文心一言统一在SSEListener中处理流式返回，通义千问与讯飞星火\智谱清言单独处理
        Boolean flag = false;
        try {
            switch (modelEnum) {

                case ZHIPU_DATA:
                    flag = sseByZhiPuData(response, sseEmitter, chatMessages, user.getUid(), chatMessage.getChatId(), conversationId, prompt, version,
                            processor, sysPrompt, ansSysPrompt);
                    break;
                case ZHIPU_DOC:
                    flag = sseByZhiPuDoc(response, sseEmitter, chatMessages, user.getUid(), chatMessage.getChatId(), conversationId, prompt, version,
                            processor, sysPrompt, ansSysPrompt);
                    break;
                case ZHIPU:
                    flag = sseByZhiPu(response, sseEmitter, chatMessages, user.getUid(), chatMessage.getChatId(), conversationId, prompt, version);
                    break;

                default:
                    throw new BusinessException("未知的模型类型，功能未接入。");
            }
            Integer status = ChatStatusEnum.SUCCESS.getValue();
            if (flag) {
                status = ChatStatusEnum.ERROR.getValue();
            }
            gptService.updateMessageStatus(conversationId, status);
        } catch (BusinessException e) {
            flag = true;
            throw e;
        } catch (Exception e) {
            e.printStackTrace();
            flag = true;
            throw new ErrorException();
        } finally {
            if (flag) {
                gptService.restoreNum(user.getId());
            }
        }
    }


    @SneakyThrows
    private Boolean sseByZhiPuData(HttpServletResponse response, SseEmitter sseEmitter, List<ChatMessageDTO> chatMessages,
                                   String uid, Long chatId, String conversationId, String prompt, String version,
                                   SqlProcessor processor, String sysPrompt, String sysPrompt2) {
        if (ValidatorUtil.isNull(zhiPuClient.getAppKey())) {
            throw new BusinessException("未加载到密钥信息");
        }

        User user = userService.getOne(new LambdaQueryWrapper<User>().eq(User::getUid, uid));

        List<ChatMessage> messages = new ArrayList<>();

        if (sysPrompt != null && !sysPrompt.isEmpty()) {
            messages.add(new ChatMessage("system", sysPrompt));
        }

        chatMessages.stream().filter(d -> !d.getRole().equals(ChatRoleEnum.SYSTEM.getValue())).forEach(v -> {
            messages.add(new ChatMessage(v.getRole(), v.getContent()));
        });
        String modelVaersion = ValidatorUtil.isNotNull(version) ? version : Constants.ModelChatGLM3TURBO;
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(modelVaersion)
                .stream(Boolean.FALSE)
                .messages(messages)
                .build();

        ModelApiResponse invokeModelApiResp = zhiPuClient.chat(chatCompletionRequest);

        if (invokeModelApiResp.isSuccess()) {


            String content = invokeModelApiResp.getData().getChoices().get(0).getMessage().getContent().toString();
            log.info("大模型返回数据:{}", content);

            String sqlJson = content.substring(content.indexOf("{"), content.lastIndexOf("}") + 1);

            JSONObject object = JSONObject.parseObject(sqlJson);

            String sql = object.getString("sqlQuery");

            log.info("返回SQL:{}", sql);

            sql = sql.replace("@tenantId", String.valueOf(user.getTenantId()));

            List<Map<String, Object>> list = processor.queryList(sql);


            log.info("查询到数据集:{}", JSONObject.toJSON(list));

            String tableData = JSONObject.toJSON(list).toString();

            List<ChatMessage> messages2 = new ArrayList<>();

            if (sysPrompt2 != null && !sysPrompt2.isEmpty()) {
                messages2.add(new ChatMessage("system", sysPrompt2));
            }
            chatMessages.stream().filter(d -> !d.getRole().equals(ChatRoleEnum.SYSTEM.getValue())).forEach(v -> {
                String qContent = v.getContent() + "\n参考数据集:" + tableData;

                messages2.add(new ChatMessage(v.getRole(), qContent));
            });
            ChatCompletionRequest chatCompletionRequest2 = ChatCompletionRequest.builder()
                    .model(modelVaersion)
                    .stream(Boolean.TRUE)
                    .messages(messages2)
                    .build();


            Boolean flag = zhiPuClient.streamChat(response, chatCompletionRequest2, chatId, conversationId, modelVaersion, processor);
            return flag;

        } else {
            log.error("智谱清言连接异常data：{}", JSON.toJSONString(invokeModelApiResp));
            return true;
        }


    }


    @SneakyThrows
    private Boolean sseByZhiPuDoc(HttpServletResponse response, SseEmitter sseEmitter, List<ChatMessageDTO> chatMessages,
                                  String uid, Long chatId, String conversationId, String prompt, String version,
                                  SqlProcessor processor, String sysPrompt, String sysPrompt2) {
        if (ValidatorUtil.isNull(zhiPuClient.getAppKey())) {
            throw new BusinessException("未加载到密钥信息");
        }

        User user = userService.getOne(new LambdaQueryWrapper<User>().eq(User::getUid, uid));

        List<ChatMessage> messages = new ArrayList<>();

        RestClient restClient = RestClient
                .builder(HttpHost.create(ES_HOST))
                .build();

        ElasticsearchTransport transport = new RestClientTransport(
                restClient, new JacksonJsonpMapper());

        ElasticsearchClient esClient = new ElasticsearchClient(transport);

        String index = "test_doc";

        SearchResponse<Map> searchResponse = esClient.search(s -> s
                        .index(index)
                        .query(q -> q
                                .match(t -> t
                                        .field("content")
                                        .query(chatMessages.get(0).getContent())
                                )
                        ),
                Map.class
        );

        StringBuffer contents = new StringBuffer();
        List<Hit<Map>> hits = searchResponse.hits().hits();
//        for (Hit<Map> hit : hits) {
//            Map map = hit.source();
//
//            contents.append(map.get("content") ).append("\n");
//        }

        if (hits != null && !hits.isEmpty()) {
            for (int i = 0; i < 2; i++) {
                Hit<Map> hit = hits.get(i);
                Map map = hit.source();
                if (hit.score() >= 1) {// 相关性评分>1才放入
                    contents.append(map.get("content")).append("\n");
                }
            }
        }

        if (sysPrompt != null && !sysPrompt.isEmpty()) {
            messages.add(new ChatMessage("system", sysPrompt));
        }


        chatMessages.stream().filter(d -> !d.getRole().equals(ChatRoleEnum.SYSTEM.getValue())).forEach(v -> {

            String question = "用户问题:" + v.getContent() + "\n";
            question += "【" + contents + "】";

            messages.add(new ChatMessage(v.getRole(), question));
        });
        String modelVaersion = ValidatorUtil.isNotNull(version) ? version : Constants.ModelChatGLM3TURBO;
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(modelVaersion)
                .stream(Boolean.TRUE)
                .messages(messages)
                .build();


        Boolean flag = zhiPuClient.streamChat(response, chatCompletionRequest, chatId, conversationId, modelVaersion, processor);
        return flag;

    }

    @SneakyThrows
    private Boolean sseByZhiPu(HttpServletResponse response, SseEmitter sseEmitter, List<ChatMessageDTO> chatMessages,
                               String uid, Long chatId, String conversationId, String prompt, String version) {
        if (ValidatorUtil.isNull(zhiPuClient.getAppKey())) {
            throw new BusinessException("未加载到密钥信息");
        }
        List<ChatMessage> messages = new ArrayList<>();
        chatMessages.stream().filter(d -> !d.getRole().equals(ChatRoleEnum.SYSTEM.getValue())).forEach(v -> {
            messages.add(new ChatMessage(v.getRole(), v.getContent()));
        });
        String modelVaersion = ValidatorUtil.isNotNull(version) ? version : Constants.ModelChatGLM3TURBO;
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(modelVaersion)
                .stream(Boolean.TRUE)
                .messages(messages)
                .build();
        Boolean flag = zhiPuClient.streamChat(response, chatCompletionRequest, chatId, conversationId, modelVaersion, null);
        return flag;
    }



}
